# %%
import matplotlib.pyplot as plt
import os
import json
import math


# %%
# k = 50
# lamb = 0.01
# # with open(os.path.join("record_torch", f"k_{k}_lambda_{lamb}", "record.json"), "r") as f:
# #     rec = json.load(f)
# with open("record.json", "r") as f:
#     rec = json.load(f)
# loss = [x[0] for x in rec]
# rmse = [math.sqrt(x[1]) for x in rec]
# plt.figure(0)
# plt.subplot(1, 2, 1)
# plt.plot(loss)
# plt.title("loss")
# plt.legend()
# plt.grid("on")
# plt.xlabel("iteration")
# plt.subplot(1, 2, 2)
# plt.plot(rmse)
# plt.title("rmse")
# plt.legend()
# plt.grid("on")
# plt.xlabel("iteration")
# plt.show()

import pandas as pd
df = pd.DataFrame(columns=[10, 50, 100, 200], index=[1, 0.1, 0.01, 0.001])
# %%
record = {}
for k in [10, 50, 100, 200]:
    for lamb in [1, 0.1, 0.01, 0.001]:
        with open(os.path.join("record", f"k_{k}_lambda_{lamb}", "record.json"), "r") as f:
            rec = json.load(f)
        rmse = rec[-1][1]
        df[k][lamb] = rmse
        print(f"k {k}, lambda {lamb}, rmse: {rmse}")
# %%
print(df)